#!/usr/bin/python
# coding=utf-8
import cookielib,time,json
from urllib2 import *

class DbCookieJar(cookielib.CookieJar):
    isload = False
    def __init__(self, policy=None):
        cookielib.CookieJar.__init__(self, policy)

    def save(self,ignore_discard=False, ignore_expires=False):
        now = time.time()
        rows = []
        for cookie in self:
            if not ignore_discard and cookie.discard:
                continue
            if not ignore_expires and cookie.is_expired(now):
                continue
            if cookie.secure: secure = "TRUE"
            else: secure = "FALSE"
            if cookie.domain.startswith("."): initial_dot = "TRUE"
            else: initial_dot = "FALSE"
            if cookie.expires is not None:
                expires = str(cookie.expires)
            else:
                expires = ""
            if cookie.value is None:
                # cookies.txt regards 'Set-Cookie: foo' as a cookie
                # with no name, whereas cookielib regards it as a
                # cookie with no value.
                name = ""
                value = cookie.name
            else:
                name = cookie.name
                value = cookie.value
            rows.append([cookie.domain, initial_dot, cookie.path, secure, expires, name, value])
        return json.dumps(rows,indent=2)

    def load(self,s,ignore_discard=False, ignore_expires=False):
        now = time.time()
        rows = json.loads(s)
        for row in rows:
            domain, domain_specified, path, secure, expires, name, value = row
            secure = (secure == "TRUE")
            domain_specified = (domain_specified == "TRUE")
            if name == "":
                # cookies.txt regards 'Set-Cookie: foo' as a cookie
                # with no name, whereas cookielib regards it as a
                # cookie with no value.
                name = value
                value = None

            initial_dot = domain.startswith(".")
            assert domain_specified == initial_dot

            discard = False
            if expires == "":
                expires = None
                discard = True

            # assume path_specified is false
            c = cookielib.Cookie(0, name, value,
                       None, False,
                       domain, domain_specified, initial_dot,
                       path, False,
                       secure,
                       expires,
                       discard,
                       None,
                       None,
                       {})
            if not ignore_discard and c.discard:
                continue
            if not ignore_expires and c.is_expired(now):
                continue
            self.set_cookie(c)
        self.isload = True

class OpenerCookies(OpenerDirector):
    cookie = DbCookieJar()
    def __init__(self):
        OpenerDirector.__init__(self)
        handler = HTTPCookieProcessor(self.cookie)
        self.build_opener(handler)
    def loadCookies(self,cookies):
        if not self.cookie.isload:
            self.cookie.load(cookies)
    def open(self, fullurl, data=None, headers={}):
        if isinstance(fullurl, basestring):
            req = Request(fullurl, data, headers)
        else:
            req = fullurl
            if data is not None:
                req.add_data(data)
                for key, value in headers.items():
                    req.add_header(key, value)
        return OpenerDirector.open(self, req)
    def build_opener(self,*handlers):
        import types
        def isclass(obj):
            return isinstance(obj, (types.ClassType, type))
        default_classes = [ProxyHandler, UnknownHandler, HTTPHandler,
                           HTTPDefaultErrorHandler, HTTPRedirectHandler,
                           FTPHandler, FileHandler, HTTPErrorProcessor]
        if hasattr(httplib, 'HTTPS'):
            default_classes.append(HTTPSHandler)
        skip = set()
        for klass in default_classes:
            for check in handlers:
                if isclass(check):
                    if issubclass(check, klass):
                        skip.add(klass)
                elif isinstance(check, klass):
                    skip.add(klass)
        for klass in skip:
            default_classes.remove(klass)

        for klass in default_classes:
            self.add_handler(klass())

        for h in handlers:
            if isclass(h):
                h = h()
            self.add_handler(h)
        return self
